#========================================
# author: changlong.zang
#   mail: zclongpop@163.com
#   date: Tue, 22 Dec 2015, 17:09:36
#========================================
import os, re
import maya.cmds as mc
#--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+

def getChildren(dagObject, dagType='transform'):
    '''
    Input a root dagNode name, find all children...
    '''
    for C in mc.listRelatives(dagObject, c=True, type=dagType, path=True) or []:
        yield C

        for c in getChildren(C, dagType):
            yield c



def convertMocapData():
    '''
    '''
    all_joints       = mc.ls(type='joint')
    mocap_joints_str = ' '.join(mc.ls('*:*', type='joint'))

    #- all joints to stand pose...
    hip = ''
    mocap_hip = ''
    for jnt in all_joints:
        mc.setAttr(jnt + '.rx', 0)
        mc.setAttr(jnt + '.ry', 0)
        mc.setAttr(jnt + '.rz', 0)

        if not re.search('hips$', jnt, re.I):
            continue

        if mc.referenceQuery(jnt, inr=True):
            mocap_hip = jnt
        else:
            hip = jnt


    posi = mc.xform(hip, q=True, ws=True, t=True)
    mc.xform(mocap_hip, ws=True, t=posi)
    mc.pointConstraint(mocap_hip, hip)


    #- match joint's rotation...
    ik_handles = list()
    for jnt in getChildren(hip, 'joint'):
        if mc.referenceQuery(jnt, inr=True):
            continue

        parent_joint = mc.listRelatives(jnt, typ='joint', p=True, path=True)
        if not parent_joint:
            continue

        typ = re.search('[a-zA-Z0-9]+$', jnt).group()
        mocap_jnt = re.search('\S+{0}'.format(typ), mocap_joints_str).group()

        ik_handle = mc.ikHandle(sj=parent_joint[0], ee=jnt, sol='ikSCsolver')[0]
        mc.pointConstraint(mocap_jnt, ik_handle)
        ik_handles.append(ik_handle)

    mc.delete(ik_handles)


    #- constraint mocap joints...
    bake_joints = list()
    for jnt in all_joints:
        if mc.referenceQuery(jnt, inr=True):
            continue

        typ = re.search('[a-zA-Z0-9]+$', jnt).group()
        mocap_jnt = re.search('\S+{0}'.format(typ), mocap_joints_str).group()

        mc.parentConstraint(mocap_jnt, jnt, mo=True)
        bake_joints.append(jnt)


    #- bake keys...
    key_frames = mc.keyframe(mocap_hip, q=True)
    start_tm = min(key_frames) - 1
    end_tm   = max(key_frames) + 1
    
    mc.playbackOptions(e=True, min=start_tm, max=end_tm, ast=start_tm, aet=end_tm)
    mc.bakeResults(bake_joints, sm=True, t=(start_tm, end_tm), hi='below', sb=1, disableImplicitControl=True, preserveOutsideKeys=True, sparseAnimCurveBake=False, removeBakedAttributeFromLayer=False, removeBakedAnimFromLayer=False, bakeOnOverrideLayer=False, minimizeRotation=False, controlPoints=False)
    mc.delete(mc.ls(typ='constraint'))



def saveNewData():
    references = mc.file(q=True, r=True)[0]
    new_name = os.path.splitext(references)[0] + '.ma'
    mc.file(references, rr=True)
    mc.file(rename=new_name)
    mc.file(save=True)
